#include <sys/types.h>          /* See NOTES */
#include <sys/socket.h>
#include <stdio.h>
#include <strings.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>

#define TIMEOUT_MS 50 // msec

int timeout_ms = TIMEOUT_MS;
int ip_num = 1;
int verbose = 0;
int socket_sniffer(char *ip, int port)
{
  int fsock;
  int result = -1;
  struct sockaddr_in target;
  struct timeval timeout;

  fsock = socket(AF_INET, SOCK_STREAM, 0);
  if (fsock < 0) {
    perror("socket create failed.");
    goto err;
  }

  timeout.tv_sec = timeout_ms / 1000;
  timeout.tv_usec = timeout_ms % 1000 * 1000;

  setsockopt(fsock, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout));
  setsockopt(fsock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

  bzero(&target, sizeof(target));
  target.sin_family = AF_INET;
  target.sin_port = htons(port);
  if (inet_pton(AF_INET, ip, &target.sin_addr) <= 0) {
    goto err;
  }

  result = connect(fsock, (struct sockaddr*)&target, sizeof(target));

err:
  if (fsock >= 0) {
    close(fsock);
  }

  return (result == 0) ? 1 : 0;
}

int check_port(char *ip, int port)
{
  int ret;
  ret = socket_sniffer(ip, port);
  //printf("sniffer ip %s , port %d, ret=%d\n", ip, port, ret);
  return ret;
}

int main (int argc, char *argv[])
{ 
  char *net = NULL;
  int port;
  int opt;
  while ((opt = getopt(argc, argv, "vt:")) != -1) {
        switch (opt) {
            case 't':
                timeout_ms = atoi(optarg);
                break;
            case 'v':
                verbose = 1;
                break;
            case '?':
                fprintf(stderr, "Usage: %s [-a ip] [-p port] \n", argv[0]);
                //exit(EXIT_FAILURE);
                exit(0);
            default:
                fprintf(stderr, "Unexpected option: %c\n", opt);
                //exit(EXIT_FAILURE);
                exit(0);
        }
    }

    if (optind + 2 != argc) {
        fprintf(stderr, "Error: Missing hostname or port\n");
        return 1;
    }
    // for (int j = 0; j < argc; j++) {
    //   printf("%d - %s\n", j, argv[j]);
    // }    
    net = argv[optind];
    port = atoi(argv[optind + 1]);


  char ip_str[20];
  int prefix_len = 0;
  uint32_t network;// = ntohl(addr.s_addr);
  uint32_t netmask;// = (prefix_len == 0 ) ? 0 : ~(1UL << (32 - prefix_len) - 1)
  uint32_t start_ip;
  struct in_addr addr;

  if (sscanf(net, "%19[^/]/%d", ip_str, &prefix_len) == 2) {
    // cidir
    if (prefix_len < 0 || prefix_len > 32) {
      perror("Invalid prefix len.");
      return -1;
    }

    ip_num = (1UL << (32 - prefix_len)) - 1;

    if (inet_pton(AF_INET, ip_str, &addr) != 1) {
        perror("Invalid ip address.");
        return -1;
    }

    network = ntohl(addr.s_addr);
    netmask = (prefix_len == 0 ) ? 0 : ~((1UL << (32 - prefix_len)) - 1);
    if (verbose) {
      printf("ip %s, ip num %d, prefix %d, network 0x%x netmask 0x%x\n", ip_str, ip_num, prefix_len, network,netmask);
    }
    
    start_ip = network & netmask;

  } else {
    strncpy(ip_str, net, sizeof(ip_str) - 1);
    if (inet_pton(AF_INET, ip_str, &addr) != 1) {
        perror("Invalid ip address.");
        return -1;
    }
    start_ip = ntohl(addr.s_addr);
  }

  //char addr[32];
  printf("Scan %s/%d :%d - timeout:%d, IPs %d, cost %ds\n", ip_str, prefix_len, port, timeout_ms, \
                                                                 ip_num, (int)(ip_num * timeout_ms / 1000.0));
  for (int i = 0; i < ip_num; i++) {
     struct in_addr start_addr;
     start_addr.s_addr = htonl(start_ip + i);
     char* start_ip_str = inet_ntoa(start_addr);

    if (verbose) {
      printf("%s-%d\n", start_ip_str, i);
    }
     if (check_port(start_ip_str, port) == 1) {
        printf("- %-16s %s\n", start_ip_str, "OPEN");
      }

  } 
  return 0;
}
